# name: 李昊旻
# hello world!
# time: 2023/4/29 10:48
import json
import numpy as np
import os
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

# 训练集和验证集的相关信息
train_yaml = "data/train.yaml"      # 训练集 YAML 文件路径
val_yaml = "data/val.yaml"          # 验证集 YAML 文件路径
train_img_dir = "data/train/images" # 训练集图片文件夹路径
val_img_dir = "data/val/images"     # 验证集图片文件夹路径

# COCO API 和评估工具的初始化
coco_gt = COCO(val_yaml)
coco_dt = coco_gt.loadRes("best_predictions.json")
coco_eval = COCOeval(coco_gt, coco_dt, "bbox")

# 准备评估数据
img_ids = sorted(coco_gt.getImgIds())
img_infos = []
gt_infos = []

for img_id in img_ids:
    img_info = coco_gt.loadImgs(img_id)[0]
    img_file = os.path.join(val_img_dir, img_info["file_name"])
    img_infos.append({"id": img_id, "file_name": img_file})
    ann_ids = coco_gt.getAnnIds(imgIds=img_id)
    ann_info = coco_gt.loadAnns(ann_ids)
    for ann in ann_info:
        del ann["mask"]
    gt_infos.extend(ann_info)


# 进行评估并计算 mAP
coco_eval.params.imgIds = img_ids
coco_eval.params.catIds = [1] # 只评估一类物体（类别 ID 为 1）
coco_eval.params.iouThrs = np.linspace(0.5, 0.95, 10) # 设定不同 IOU 阈值
coco_eval.evaluate()
coco_eval.accumulate()
coco_eval.summarize()
mAP = coco_eval.stats[0] # 获取 mAP 值